import torch
import torch.nn as nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 定义第一卷积层：输入通道1，输出通道16，3x3卷积核，保持输入尺寸
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        # 定义最大池化层：2x2窗口，步长2，将特征图尺寸减半
        self.pool = nn.MaxPool2d(2, 2)
        # 定义第二卷积层：输入通道16，输出通道32，3x3卷积核，保持输入尺寸
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        # 定义第一全连接层：将32个7x7特征图展平后映射到128维
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        # 定义输出全连接层：将128维特征映射到10类输出
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 第一卷积块：卷积->ReLU->池化
        # 输入：(B,1,28,28) -> Conv2d -> (B,16,28,28) -> MaxPool -> (B,16,14,14)
        x = self.pool(torch.relu(self.conv1(x)))
        # 第二卷积块：卷积->ReLU->池化
        # 输入：(B,16,14,14) -> Conv2d -> (B,32,14,14) -> MaxPool -> (B,32,7,7)
        x = self.pool(torch.relu(self.conv2(x)))
        # 特征展平：将32个7x7特征图展平为一维向量
        # 形状变换：(B,32,7,7) -> view -> (B, 32*7*7)
        x = x.view(-1, 32 * 7 * 7)
        # 全连接层处理：128维隐藏层 + ReLU激活
        x = torch.relu(self.fc1(x))
        # 输出层：无激活函数的线性变换
        x = self.fc2(x)
        return x
